from kgdlg.modules.Encoder import BasicEncoder
from kgdlg.modules.Decoder import AttnDecoder,KeyAttnDecoder
from kgdlg.modules.Embedding import Embedding,ClusterEmbedding
from kgdlg.Model import NMTModel,similar_PriorNet,similar_RecogNet,MultiGen,diverse_PriorNet,diverse_RecogNet
import torch
import torch.nn as nn
import kgdlg.IO as IO



def create_encoder(opt, embedding):
    
    rnn_type = opt.rnn_type
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout
    bidirectional = opt.bidirectional

    encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)

    return encoder

def create_decoder(opt,embedding):

    decoder_type = opt.decoder_type
    rnn_type = opt.rnn_type  
    atten_model = opt.atten_model
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout 

    if decoder_type == 'AttnDecoder':
        print("当前的decoder为AttnDecoder")
        decoder = KeyAttnDecoder(rnn_type,
                                embedding,
                                atten_model,
                                input_size,
                                hidden_size,
                                num_layers,
                                dropout)


    return decoder

def create_generator(input_size, output_size):
    generator = nn.Sequential(
        nn.Linear(input_size, output_size),
        nn.LogSoftmax(dim=-1))
    return generator


def create_similar_prior_model(opt, fields):
    vocab_size = len(fields['tgt'].vocab)
    padding_idx = fields['tgt'].vocab.stoi[IO.PAD_WORD]
    embedding_size = opt.embedding_size
    embedding = Embedding(vocab_size,embedding_size, padding_idx)

    rnn_type = opt.rnn_type
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout
    bidirectional = opt.bidirectional
    encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)

    similar_prior_net = similar_PriorNet(encoder)

    return similar_prior_net

def create_diverse_PR_model(opt, fields):
    vocab_size = len(fields['tgt'].vocab)
    padding_idx = fields['tgt'].vocab.stoi[IO.PAD_WORD]
    embedding_size = opt.embedding_size
    embedding = Embedding(vocab_size,embedding_size, padding_idx)

    rnn_type = opt.rnn_type
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout
    bidirectional = opt.bidirectional
    encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)

    similar_prior_net = similar_PriorNet(encoder)

    return similar_prior_net

def create_similar_recog_model(opt, fields):
    vocab_size = len(fields['tgt'].vocab)
    padding_idx = fields['tgt'].vocab.stoi[IO.PAD_WORD]
    embedding_size = opt.embedding_size
    embedding = Embedding(vocab_size,embedding_size, padding_idx)

    rnn_type = opt.rnn_type
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout
    bidirectional = opt.bidirectional
    src_encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)
    tgt_encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)

    similar_recog_model = similar_RecogNet(src_encoder, tgt_encoder)

    return similar_recog_model


def create_joint_model(opt, fields):
    print("当前的训练模式：", opt.train_mode)

    vocab_size = len(fields['tgt'].vocab)
    padding_idx = fields['tgt'].vocab.stoi[IO.PAD_WORD]
    embedding_size = opt.embedding_size
    embedding = Embedding(vocab_size,embedding_size, padding_idx)

    similar_prior_model = create_similar_prior_model(opt, fields)
    similar_recog_model = create_similar_recog_model(opt, fields)
    diverse_prior_net = create_diverse_prior_model(opt, fields)
    diverse_recog_net = create_diverse_recog_model(opt, fields)
    encoder = create_encoder(opt, embedding)
    decoder = create_decoder(opt,embedding)
    generator = create_generator(opt.hidden_size, vocab_size)
    mask_set = set(list(range(vocab_size)))
    if opt.train_mode == 200:
        diverse_PR_net = create_similar_prior_model(opt, fields)
        joint_model = MultiGen(similar_prior_model, similar_recog_model, diverse_prior_net, diverse_recog_net,
                               encoder, decoder, generator,
                               mask_set, opt,diverse_PR_net)
    else:
        joint_model = MultiGen(similar_prior_model, similar_recog_model, diverse_prior_net, diverse_recog_net,
                                encoder, decoder, generator,
                                 mask_set,opt)

    return joint_model


def create_diverse_prior_model(opt, fields):
    vocab_size = len(fields['tgt'].vocab)
    padding_idx = fields['tgt'].vocab.stoi[IO.PAD_WORD]
    embedding_size = opt.embedding_size
    embedding = Embedding(vocab_size,embedding_size, padding_idx)

    rnn_type = opt.rnn_type
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout
    bidirectional = opt.bidirectional
    encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)
    diverse_prior_net = diverse_PriorNet(encoder)

    return diverse_prior_net

def create_diverse_recog_model(opt, fields):
    vocab_size = len(fields['tgt'].vocab)
    padding_idx = fields['tgt'].vocab.stoi[IO.PAD_WORD]
    embedding_size = opt.embedding_size
    embedding = Embedding(vocab_size,embedding_size, padding_idx)

    rnn_type = opt.rnn_type
    input_size = opt.embedding_size
    hidden_size = opt.hidden_size
    num_layers = opt.num_layers
    dropout = opt.dropout
    bidirectional = opt.bidirectional
    src_encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)
    tgt_encoder = BasicEncoder(rnn_type,
                        embedding,
                        input_size,
                        hidden_size,
                        num_layers,
                        dropout,
                        bidirectional)
    diverse_recog_net = diverse_RecogNet(src_encoder, tgt_encoder)

    return diverse_recog_net

